﻿#include "process_query.h"

/**
 * @brief 查询前的准备工作
 */
void process_query::prepare()
{
    set_security();
    create_instance();
    connect_server();
    set_proxy_blanket();
}

/**
 * @brief 发起WMI查询请求
 *
 * @return 查询到的进程信息
 */
std::vector<process_query_info> process_query::execute_query()
{
    com_interface<IEnumWbemClassObject> enumerator;
    HRESULT hResult = service->ExecQuery(CComBSTR(L"WQL"),
        CComBSTR(L"SELECT * from win32_process"),
        WBEM_FLAG_FORWARD_ONLY | WBEM_FLAG_RETURN_IMMEDIATELY,
        nullptr,
        enumerator.address());
    check_result(hResult, L"ExecQuery");
    return fetch_results(std::move(enumerator));
}

/**
 * @brief 设置COM安全信息
 */
void process_query::set_security()
{
    HRESULT hResult = CoInitializeSecurity(nullptr,
        -1,
        nullptr,
        nullptr,
        RPC_C_AUTHN_LEVEL_DEFAULT,
        RPC_C_IMP_LEVEL_IMPERSONATE,
        nullptr,
        EOAC_NONE,
        nullptr);
    check_result(hResult, L"CoInitializeSecurity");
}

/**
 * @brief 创建进程内COM服务器
 */
void process_query::create_instance()
{
    // 创建进程内COM服务器
    HRESULT hResult = CoCreateInstance(CLSID_WbemLocator,
        nullptr,
        CLSCTX_INPROC_SERVER,
        IID_IWbemLocator,
        (LPVOID*)locator.address());    // 需要强制转型
    check_result(hResult, L"CoCreateInstance");
}

/**
 * @brief 连接WMI命名空间
 */
void process_query::connect_server()
{
    // 不能直接传递L"Root\\CIMv2"，const BSTR和const WCHAR*不一样
    HRESULT hResult = locator->ConnectServer(CComBSTR(L"Root\\CIMv2"),
        nullptr,
        nullptr,
        nullptr,
        0,
        nullptr,
        nullptr,
        service.address());
    check_result(hResult, L"ConnectServer");
}

/**
 * @brief 设置WMI连接的安全等级
 */
void process_query::set_proxy_blanket()
{
    HRESULT hResult = CoSetProxyBlanket(service.get(),
        RPC_C_AUTHN_WINNT,
        RPC_C_AUTHZ_NONE,
        nullptr,
        RPC_C_AUTHN_LEVEL_CALL,
        RPC_C_IMP_LEVEL_IMPERSONATE,
        nullptr,
        EOAC_NONE);
    check_result(hResult, L"CoSetProxyBlanket");
}

/**
 * @brief 提取查询结果
 *
 * @param enumerator 枚举对象
 * @return 查询结果
 */
std::vector<process_query_info> process_query::fetch_results(
    com_interface<IEnumWbemClassObject> enumerator)
{
    std::vector<process_query_info> vec;
    while (get_next_process_info(enumerator.get(), vec))
    {
        // do nothing
    }
    return vec;
}

/**
 * @brief 获取下一个进程信息
 *
 * @param enumerator 枚举对象
 * @param vec 存放信息的容器
 * @return 获取成功返回TRUE，否则返回FALSE
 */
BOOL process_query::get_next_process_info(
    IEnumWbemClassObject* enumerator,
    std::vector<process_query_info>& vec)
{
    com_interface<IWbemClassObject> obj;
    ULONG ret = 0;
    // Next的返回值不能直接反馈是否操作成功，因此这里不进行校验
    enumerator->Next(WBEM_INFINITE, 1, obj.address(), &ret);
    if (ret == 0)
    {
        return FALSE;  // 已获取全部数据
    }

    VARIANT pid;        // 进程id
    VARIANT name;       // 进程名称
    VARIANT cmd_line;   // 进程命令行

    VariantInit(&pid);
    VariantInit(&name);
    VariantInit(&cmd_line);

    // 这里忽略错误检查
    obj->Get(L"ProcessId", 0, &pid, nullptr, nullptr);
    obj->Get(L"Name", 0, &name, nullptr, nullptr);
    obj->Get(L"CommandLine", 0, &cmd_line, nullptr, nullptr);

    vec.emplace_back(pid.ulVal,
        name.bstrVal ? name.bstrVal : L"",
        cmd_line.bstrVal ? cmd_line.bstrVal : L"");

    VariantClear(&pid);
    VariantClear(&name);
    VariantClear(&cmd_line);

    return TRUE;
}
